import unittest

from tinydb.sql.logical_operator import Condition
from tinydb.sql.parser.ast import BinaryOperation, Constant, Identifier, JoinType
from tinydb.common.fabric import TableColumn
from tinydb.executor.operator.physical_operator import NestedLoopJoin, TableScan, cast_tuple_pair_to_values

def construct_condition(sign, column, value):
    b = BinaryOperation(op=sign, args=(
        Identifier(column), Constant(value)))
    return Condition(b)


def construct_join_condition(sign, column1, column2):
    b = BinaryOperation(op=sign, args=(
        Identifier(column1), Identifier(column2)))
    return Condition(b)


class TestExecutor(unittest.TestCase):
    def test_cast_tuple_pair_to_values(self):
        columns = (
            TableColumn('t1', 'id'),
            TableColumn('t1', 'name'),
            TableColumn('t1', 'age')
        )
        excepted = {TableColumn('t1', 'id'): 1, TableColumn('t1', 'name'): 'a', TableColumn('t1', 'age'): 10}
        self.assertEqual(
            cast_tuple_pair_to_values(columns, (1, 'a', 10)),
            excepted)
        
    def test_nested_loop_join(self):
        # inner join
        opt = NestedLoopJoin(
            join_type=JoinType.INNER_JOIN,
            left_table_name='t1',
            right_table_name='t2',
            join_condition=construct_join_condition('=', 't1.id', 't2.id')
        )
        opt.add_child(TableScan(table_name='t1'))
        opt.add_child(TableScan(table_name='t2'))
        opt.open()
        expected = [(1, '小明', 1, '明小', '北京'), (2, '小红', 2, '红柒', '上海'), (3, '小刚', 3, '刚营', '广州'), (4, '小丽', 4, '丽江', '云南')]
        self.assertEqual(list(opt.next()), expected)
        opt.close()

        # cross join
        opt = NestedLoopJoin(
            join_type=JoinType.CROSS_JOIN,
            left_table_name='t1',
            right_table_name='t2',
            join_condition=construct_join_condition('=', 't1.id', 't2.id')
        )
        opt.add_child(TableScan(table_name='t1'))
        opt.add_child(TableScan(table_name='t2'))
        opt.open()
        expected = [(1, '小明', 1, '明小', '北京'),
                (1, '小明', 2, '红柒', '上海'),
                (1, '小明', 3, '刚营', '广州'),
                (1, '小明', 4, '丽江', '云南'),
                (2, '小红', 1, '明小', '北京'),
                (2, '小红', 2, '红柒', '上海'),
                (2, '小红', 3, '刚营', '广州'),
                (2, '小红', 4, '丽江', '云南'),
                (3, '小刚', 1, '明小', '北京'),
                (3, '小刚', 2, '红柒', '上海'),
                (3, '小刚', 3, '刚营', '广州'),
                (3, '小刚', 4, '丽江', '云南'),
                (4, '小丽', 1, '明小', '北京'),
                (4, '小丽', 2, '红柒', '上海'),
                (4, '小丽', 3, '刚营', '广州'),
                (4, '小丽', 4, '丽江', '云南')]
        self.assertEqual(list(opt.next()), expected)
        opt.close()

        # left join
        opt = NestedLoopJoin(
            join_type=JoinType.LEFT_JOIN,
            left_table_name='t1',
            right_table_name='t2',
            join_condition=construct_join_condition('=', 't1.id', 't2.id')
        )
        opt.add_child(TableScan(table_name='t1'))
        opt.add_child(TableScan(table_name='t2'))
        opt.open()
        expected = [(1, '小明', 1, '明小', '北京'), (2, '小红', 2, '红柒', '上海'), (3, '小刚', 3, '刚营', '广州'), (4, '小丽', 4, '丽江', '云南')]
        self.assertEqual(list(opt.next()), expected)
        opt.close()

        # right join
        opt = NestedLoopJoin(
            join_type=JoinType.RIGHT_JOIN,
            left_table_name='t1',
            right_table_name='t2',
            join_condition=construct_join_condition('=', 't1.id', 't2.id')
        )
        opt.add_child(TableScan(table_name='t1'))
        opt.add_child(TableScan(table_name='t2'))
        opt.open()
        expected = [(1, '小明', 1, '明小', '北京'), (2, '小红', 2, '红柒', '上海'), (3, '小刚', 3, '刚营', '广州'), (4, '小丽', 4, '丽江', '云南')]
        self.assertEqual(list(opt.next()), expected)
        opt.close()

        # inner join
        opt = NestedLoopJoin(
            join_type=JoinType.INNER_JOIN,
            left_table_name='t1',
            right_table_name='t3',
            join_condition=construct_join_condition('=', 't1.id', 't3.id')
        )
        opt.add_child(TableScan(table_name='t1'))
        opt.add_child(TableScan(table_name='t3'))
        opt.open()
        expected = [(1, '小明', 1, '北京', '朝阳'), (2, '小红', 2, '上海', '浦东'), (3, '小刚', 3, '广州', '天河')]
        self.assertEqual(list(opt.next()), expected)
        opt.close()

        # left join
        opt = NestedLoopJoin(
            join_type=JoinType.LEFT_JOIN,
            left_table_name='t1',
            right_table_name='t3',
            join_condition=construct_join_condition('=', 't1.id', 't3.id')
        )
        opt.add_child(TableScan(table_name='t1'))
        opt.add_child(TableScan(table_name='t3'))
        opt.open()
        expected = [(1, '小明', 1, '北京', '朝阳'), (2, '小红', 2, '上海', '浦东'), (3, '小刚', 3, '广州', '天河'), (4, '小丽', None, None, None)]
        self.assertEqual(list(opt.next()), expected)
        opt.close()

        # right join
        opt = NestedLoopJoin(
            join_type=JoinType.RIGHT_JOIN,
            left_table_name='t1',
            right_table_name='t3',
            join_condition=construct_join_condition('=', 't1.id', 't3.id')
        )
        opt.add_child(TableScan(table_name='t1'))
        opt.add_child(TableScan(table_name='t3'))
        opt.open()
        expected = [(1, '小明', 1, '北京', '朝阳'), (2, '小红', 2, '上海', '浦东'), (3, '小刚', 3, '广州', '天河')]
        self.assertEqual(list(opt.next()), expected)
        opt.close()

        # cross join
        opt = NestedLoopJoin(
            join_type=JoinType.CROSS_JOIN,
            left_table_name='t1',
            right_table_name='t3',
            join_condition=construct_join_condition('=', 't1.id', 't3.id')
        )
        opt.add_child(TableScan(table_name='t1'))
        opt.add_child(TableScan(table_name='t3'))
        opt.open()
        expected = [(1, '小明', 1, '北京', '朝阳'), (1, '小明', 2, '上海', '浦东'), (1, '小明', 3, '广州', '天河'), (2, '小红', 1, '北京', '朝阳'), (2, '小红', 2, '上海', '浦东'), (2, '小红', 3, '广州', '天河'), (3, '小刚', 1, '北京', '朝阳'), (3, '小刚', 2, '上海', '浦东'), (3, '小刚', 3, '广州', '天河'), (4, '小丽', 1, '北京', '朝阳'), (4, '小丽', 2, '上海', '浦东'), (4, '小丽', 3, '广州', '天河')]
        self.assertEqual(list(opt.next()), expected)
        opt.close()

        # full join
        opt = NestedLoopJoin(
            join_type=JoinType.FULL_JOIN,
            left_table_name='t1',
            right_table_name='t3',
            join_condition=construct_join_condition('=', 't1.id', 't3.id')
        )
        opt.add_child(TableScan(table_name='t1'))
        opt.add_child(TableScan(table_name='t3'))
        opt.open()
        expected = [(1, '小明', 1, '北京', '朝阳'), (2, '小红', 2, '上海', '浦东'), (3, '小刚', 3, '广州', '天河'), (4, '小丽', None, None, None)]
        self.assertEqual(list(opt.next()), expected)
        opt.close()

        opt = NestedLoopJoin(
            join_type=JoinType.FULL_JOIN,
            left_table_name='t1',
            right_table_name='t4',
            join_condition=construct_join_condition('=', 't1.id', 't4.id')
        )
        opt.add_child(TableScan(table_name='t1'))
        opt.add_child(TableScan(table_name='t4'))
        opt.open()
        expected = [(1, '小明', None, None, None), (None, None, 5, 'aws', 'sales'), (2, '小红', 2, 'google', 'beckend'), (None, None, 5, 'aws', 'sales'), (3, '小刚', 3, 'meta', 'frontend'), (None, None, 5, 'aws', 'sales'), (4, '小丽', None, None, None), (None, None, 5, 'aws', 'sales')]
        self.assertEqual(list(opt.next()), expected)
        opt.close()



if __name__ == '__main__':
    unittest.main()
